import scipy.io as scio
from scipy.misc import imread,imshow
import tensorflow as tf
import numpy as np
from PIL import Image
import scipy.ndimage.filters as scfil

data=scio.loadmat("/home/chen/DCNN_caffe-Tensorflow/test1.mat")


def weight_variable(name):
    return data[name]
    '''
    initial=data[name]
    x = tf.sqrt(tf.cast(tf.shape(initial)[1], tf.float32))
    x=tf.cast(x,tf.int32)
    y = tf.shape(initial)[2]
    initial = tf.reshape(initial, [-1,x,x,y])
    initial = tf.transpose(initial, [2,1,0,3])
    initial=tf.cast(tf.reshape(initial,shape),tf.float32)
    return tf.Variable(initial)
    '''


def bias_variable(name):
    #initial=tf.constant(data[name],shape)
     initial=data[name]
     #initial=tf.cast(tf.reshape(initial,shape),tf.float32)
     return initial


def conv2d(x,W):
    return tf.nn.conv2d(x,W,strides=[1,1,1,1],padding='SAME')
#[1,x_movement,y_movement,1], padding=SAME/VALID


def max_pool_2x2(x):
    return tf.nn.max_pool(x,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')
#[1,x_movement,y_movement,1], padding=SAME/VALID





def getImage(image_path,knn_path,closedform_path):
    ClosedForm_image=tf.cast(imread(closedform_path),tf.float32)
    KNN_image=tf.cast(imread(knn_path),tf.float32)
    x_image=tf.cast(imread(image_path),tf.float32)
    with tf.Session() as sess:
         s=x_image[:,:,0]*x_image[:,:,0]+x_image[:,:,1]*x_image[:,:,1]+x_image[:,:,2]*x_image[:,:,2]
         s=sess.run(s)
    s=tf.pack(s)
    x_image_r= tf.div(x_image[:,:,0],s)
    x_image_g= tf.div(x_image[:,:,1],s)
    x_image_b= tf.div(x_image[:,:,2],s)
    with tf.Session() as sess:
         ClosedForm_image=sess.run(ClosedForm_image)
         KNN_image=sess.run(KNN_image)
         # print(np.shape(sess.run(x_image_r)))
         #x_image=[sess.run(x_image_r),sess.run(x_image_g),sess.run(x_image_b)]
         x_image=[sess.run(x_image_r),sess.run(x_image_g),sess.run(x_image_b),ClosedForm_image, KNN_image]
         x_image = tf.pack(x_image,axis=2)
         x_image=sess.run(x_image)-0.5
         return x_image





def Conv1(Image,Kernel,biases):
    x = np.sqrt(np.shape(Kernel)[1])
    conv_patchsize=np.int32(x)
    conv_filters = np.shape(Kernel)[2]
    hei=np.shape(Image)[0]
    wid=np.shape(Image)[1]
    ch=np.shape(Image)[2]
    Kernel = np.reshape(Kernel, (ch,conv_patchsize,conv_patchsize, conv_filters))
    Kernel=np.float32(Kernel)
    hei=np.int32(hei)
    wid=np.int32(wid)
    conv_filters=np.int32(conv_filters)
    conv_data=np.zeros((hei, wid,  conv_filters))
    for i in range(0,conv_filters-1):
        for j in range(0,ch-1):
            conv_subfilter=Kernel[j,:,:,i]
            conv_subfilter =np.reshape(conv_subfilter, (conv_patchsize, conv_patchsize));
            sub_image=Image[:,:,j]
            sub_image=np.reshape(sub_image,(hei,wid))
            filter_data=np.reshape(scfil.convolve(sub_image,conv_subfilter),(hei,wid))
            conv_seq_i=np.add(np.float32(conv_data[:,:,i]),np.float32(filter_data))
        conv_array=np.add(conv_seq_i,biases[i])
        conv_array=np.array(conv_array)
        conv_array[conv_array<0]=0
        conv_data[:,:,i]=conv_array
    return conv_data  




def Conv2(Image,Kernel,biases):
    x = np.sqrt(np.shape(Kernel)[1])
    conv_patchsize=np.int32(x)
    conv_filters = np.shape(Kernel)[2]
    hei=np.shape(Image)[0]
    wid=np.shape(Image)[1]
    ch=np.shape(Image)[2]
    Kernel=np.float32(Kernel)
    hei=np.int32(hei)
    wid=np.int32(wid)
    conv_filters=np.int32(conv_filters)
    conv_data=np.zeros((hei, wid,  conv_filters))
    for i in range(0,conv_filters-1):
        Kernel_i=np.reshape(Kernel[:,:,i],(1,1,ch))
        repmat=np.tile(Kernel_i,(hei,wid,1))
        repmat=np.array(repmat)
        Image=np.array(Image)
        conv_data[:,:,i]=np.sum((repmat*Image),2)
        conv_array=np.add(conv_data[:,:,i],biases[i])
        conv_array=np.array(conv_array)
        conv_array[conv_array<0]=0
        conv_data[:,:,i]=conv_array
    return conv_data





def Conv3(Image,Kernel,biases):
    x = np.sqrt(np.shape(Kernel)[1])
    conv_patchsize=np.int32(x)
    hei=np.shape(Image)[0]
    wid=np.shape(Image)[1]
    hei=np.int32(hei)
    wid=np.int32(wid)
    ch=np.shape(Image)[2]
    conv_data=np.zeros((hei, wid))
    for i in range (0,ch-1):
        conv_subfilter=Kernel[i,:]
        conv_subfilter =np.reshape(conv_subfilter, (conv_patchsize, conv_patchsize));
        filter_data=scfil.convolve(Image[:,:,i],conv_subfilter)
        conv_data=np.add(conv_data,filter_data)
    return conv_data  



    
    
    




#x_image=tf.reshape(xs,[-1,28,28,1])    #[n_samples,width,height,channel]
image_path='/home/chen/DCNN_caffe-Tensorflow/image.png'
closedform_path='/home/chen/DCNN_caffe-Tensorflow/ClosedForm.png'
knn_path='/home/chen/DCNN_caffe-Tensorflow/KNN.png'
x_image=getImage(image_path,closedform_path,knn_path)
x_image=x_image-0.5
trimap=imread('/home/chen/DCNN_caffe-Tensorflow/trimap.png')




#================================================


sess=tf.Session()

#conv1 layer
W_conv1 = weight_variable('weights_conv1')    #9x9 patch, in size 5, out size 64
#W_conv1=W_conv1[:,:,0:3,:]
b_conv1=bias_variable('biases_conv1')
h_conv1=tf.nn.relu(Conv1(x_image,W_conv1,b_conv1)) #output size 
#h_conv1=conv2d(x_image,W_conv1)+b_conv1
h_conv1=tf.nn.dropout(h_conv1,0.7)
h_conv1=sess.run(h_conv1)



#conv2_1 layer
W_conv2_1 = weight_variable('weights_conv2_1')    #1x1 patch, in size 64, out size 64
b_conv2_1=bias_variable('biases_conv2_1')
h_conv2_1=tf.nn.relu(Conv2(h_conv1,W_conv2_1,b_conv2_1)) 
h_conv2_1=sess.run(h_conv2_1)


#conv2_2 layer
W_conv2_2 = weight_variable('weights_conv2_2')    #1x1 patch, in size 64, out size 64
b_conv2_2=bias_variable('biases_conv2_2')
h_conv2_2=tf.nn.relu(Conv2(h_conv2_1,W_conv2_2,b_conv2_2))
h_conv2_2=sess.run(h_conv2_2) 

 
#conv2_3 layer
W_conv2_3 = weight_variable('weights_conv2_3')    #1x1 patch, in size 64, out size 64
b_conv2_3=bias_variable('biases_conv2_3')
h_conv2_3=tf.nn.relu(Conv2(h_conv2_2,W_conv2_3,b_conv2_3))
h_conv2_3=sess.run(h_conv2_3) 



#conv2_4 layer
W_conv2_4 = weight_variable('weights_conv2_4')    #1x1 patch, in size 64, out size 32
b_conv2_4=bias_variable('biases_conv2_4')
h_conv2_4=tf.nn.relu(Conv2(h_conv2_3,W_conv2_4,b_conv2_4)) 
h_conv2_4=sess.run(h_conv2_4)


#conv3 layer
W_conv3 = weight_variable('weights_conv3')     #5x5 patch, in size 32, out size 1
b_conv3=bias_variable('biases_conv3')
h_conv3=Conv3(h_conv2_4,W_conv3,b_conv3) #output size 
#h_conv1=conv2d(x_image,W_conv1)+b_conv1



prediction=h_conv3+0.5
prediction=np.array(prediction,np.float32)
#prediction=np.dot(prediction,255)


prediction[prediction<0]=0
prediction[prediction>255]=255
x=np.shape(prediction)[0]
y=np.shape(prediction)[1]


for i in range (0,x-1):
    for j in range (0,y-1):
        if (trimap[i,j].max()<1):
            prediction[i,j]=0.0
        if (trimap[i,j].min()>254):
            prediction[i,j]=255.0


closedform=imread(closedform_path)
knn=imread(knn_path)
for i in range (0,x-1):
    for j in range (0,y-1):
        if (np.abs(closedform[i,j]-knn[i,j])<30):
            prediction[i,j]=closedform[i,j]

prediciton=np.array(prediction)
prediction=np.uint8(prediction)
prediction=Image.fromarray(prediction)
prediction.convert("L")
prediction.save('/home/chen/learn_tensorflow/prediciton.png')

imshow(prediction)


